"""
技术指标计算模块（核心）
从stock_analyzer.py迁移，保持API不变
"""

import pandas as pd
import numpy as np
from config import current_config as config


class TechnicalIndicators:
    """技术指标计算类（保持向后兼容）"""
    
    @staticmethod
    def calculate_ma(df, periods=None):
        """计算移动平均线（金融标准）"""
        if periods is None:
            periods = config.MA_PERIODS
        
        for period in periods:
            # 使用 min_periods=period，确保每条均线使用完整周期数据
            # 这是金融图表的标准做法：MA5从第5日开始，MA10从第10日开始，MA20从第20日开始
            # 避免在数据起始处均线重合的问题
            df[f'MA{period}'] = df['收盘'].rolling(window=period, min_periods=period).mean()
        return df
    
    @staticmethod
    def calculate_ema(df, period=12):
        """计算指数移动平均线"""
        return df['收盘'].ewm(span=period, adjust=False).mean()
    
    @staticmethod
    def calculate_macd(df):
        """计算MACD指标"""
        ema12 = df['收盘'].ewm(span=config.MACD_FAST, adjust=False).mean()
        ema26 = df['收盘'].ewm(span=config.MACD_SLOW, adjust=False).mean()
        
        df['DIF'] = ema12 - ema26
        df['DEA'] = df['DIF'].ewm(span=config.MACD_SIGNAL, adjust=False).mean()
        df['MACD'] = 2 * (df['DIF'] - df['DEA'])
        
        return df
    
    @staticmethod
    def calculate_kdj(df, n=None):
        """计算KDJ指标（金融标准）"""
        if n is None:
            n = config.KDJ_PERIOD
        
        # 明确指定 min_periods，确保使用完整周期数据
        low_list = df['最低'].rolling(window=n, min_periods=n).min()
        high_list = df['最高'].rolling(window=n, min_periods=n).max()
        
        rsv = (df['收盘'] - low_list) / (high_list - low_list) * 100
        
        df['K'] = rsv.ewm(com=2, adjust=False).mean()
        df['D'] = df['K'].ewm(com=2, adjust=False).mean()
        df['J'] = 3 * df['K'] - 2 * df['D']
        
        return df
    
    @staticmethod
    def calculate_rsi(df, period=14):
        """计算RSI相对强弱指标（金融标准）"""
        delta = df['收盘'].diff()
        # 明确指定 min_periods，确保使用完整周期数据
        gain = (delta.where(delta > 0, 0)).rolling(window=period, min_periods=period).mean()
        loss = (-delta.where(delta < 0, 0)).rolling(window=period, min_periods=period).mean()
        
        rs = gain / loss
        df[f'RSI{period}'] = 100 - (100 / (1 + rs))
        
        return df
    
    @staticmethod
    def calculate_boll(df, period=None, std_multiplier=None):
        """计算布林带（金融标准）"""
        if period is None:
            period = config.BOLL_PERIOD
        if std_multiplier is None:
            std_multiplier = config.BOLL_STD
        
        # 使用完整周期数据，确保布林带的准确性
        # 中轨是MA，上下轨基于标准差
        df['BOLL_MID'] = df['收盘'].rolling(window=period, min_periods=period).mean()
        df['BOLL_STD'] = df['收盘'].rolling(window=period, min_periods=period).std()
        df['BOLL_UPPER'] = df['BOLL_MID'] + std_multiplier * df['BOLL_STD']
        df['BOLL_LOWER'] = df['BOLL_MID'] - std_multiplier * df['BOLL_STD']
        
        return df
    
    @staticmethod
    def calculate_volume_ma(df, periods=[5, 10, 20]):
        """计算成交量均线（金融标准）"""
        for period in periods:
            # 使用完整周期数据，与价格均线保持一致
            # VOL_MA5从第5日开始，VOL_MA10从第10日开始
            df[f'VOL_MA{period}'] = df['成交量'].rolling(window=period, min_periods=period).mean()
        return df
    
    @staticmethod
    def calculate_wr(df, period=14):
        """计算威廉指标(Williams %R)（金融标准）"""
        # 明确指定 min_periods，确保使用完整周期数据
        high_max = df['最高'].rolling(window=period, min_periods=period).max()
        low_min = df['最低'].rolling(window=period, min_periods=period).min()
        
        df[f'WR{period}'] = -100 * (high_max - df['收盘']) / (high_max - low_min)
        return df
    
    @staticmethod
    def calculate_cci(df, period=14):
        """计算CCI指标(Commodity Channel Index)（金融标准）"""
        tp = (df['最高'] + df['最低'] + df['收盘']) / 3
        # 明确指定 min_periods，确保使用完整周期数据
        ma = tp.rolling(window=period, min_periods=period).mean()
        md = tp.rolling(window=period, min_periods=period).apply(lambda x: abs(x - x.mean()).mean())
        
        df['CCI'] = (tp - ma) / (0.015 * md)
        return df
    
    @staticmethod
    def calculate_atr(df, period=14):
        """计算ATR指标(Average True Range) - 真实波动幅度（金融标准）"""
        high_low = df['最高'] - df['最低']
        high_close = abs(df['最高'] - df['收盘'].shift())
        low_close = abs(df['最低'] - df['收盘'].shift())
        
        tr = pd.concat([high_low, high_close, low_close], axis=1).max(axis=1)
        # 明确指定 min_periods，确保使用完整周期数据
        df['ATR'] = tr.rolling(window=period, min_periods=period).mean()
        return df
    
    @staticmethod
    def calculate_obv(df):
        """计算OBV指标(On Balance Volume) - 能量潮"""
        obv = [0]
        for i in range(1, len(df)):
            if df['收盘'].iloc[i] > df['收盘'].iloc[i-1]:
                obv.append(obv[-1] + df['成交量'].iloc[i])
            elif df['收盘'].iloc[i] < df['收盘'].iloc[i-1]:
                obv.append(obv[-1] - df['成交量'].iloc[i])
            else:
                obv.append(obv[-1])
        
        df['OBV'] = obv
        return df
    
    @staticmethod
    def calculate_dmi(df, period=14):
        """计算DMI指标(Directional Movement Index) - 趋向指标（金融标准）"""
        # 计算+DM和-DM
        high_diff = df['最高'].diff()
        low_diff = -df['最低'].diff()
        
        plus_dm = high_diff.where((high_diff > low_diff) & (high_diff > 0), 0)
        minus_dm = low_diff.where((low_diff > high_diff) & (low_diff > 0), 0)
        
        # 计算TR
        high_low = df['最高'] - df['最低']
        high_close = abs(df['最高'] - df['收盘'].shift())
        low_close = abs(df['最低'] - df['收盘'].shift())
        tr = pd.concat([high_low, high_close, low_close], axis=1).max(axis=1)
        
        # 平滑处理 - 明确指定 min_periods，确保使用完整周期数据
        atr = tr.rolling(window=period, min_periods=period).sum()
        plus_di = 100 * plus_dm.rolling(window=period, min_periods=period).sum() / atr
        minus_di = 100 * minus_dm.rolling(window=period, min_periods=period).sum() / atr
        
        # 计算ADX
        dx = 100 * abs(plus_di - minus_di) / (plus_di + minus_di)
        adx = dx.rolling(window=period, min_periods=period).mean()
        
        df['PDI'] = plus_di
        df['MDI'] = minus_di
        df['ADX'] = adx
        
        return df
    
    @staticmethod
    def calculate_bias(df, periods=[6, 12, 24]):
        """计算BIAS乖离率（金融标准）"""
        for period in periods:
            if len(df) >= period:
                # 明确指定 min_periods，确保使用完整周期数据
                ma = df['收盘'].rolling(window=period, min_periods=period).mean()
                df[f'BIAS{period}'] = (df['收盘'] - ma) / ma * 100
        return df
    
    @staticmethod
    def calculate_roc(df, period=12):
        """计算ROC指标(Price Rate of Change) - 变动率指标"""
        df['ROC'] = (df['收盘'] - df['收盘'].shift(period)) / df['收盘'].shift(period) * 100
        return df
    
    @staticmethod
    def calculate_psy(df, period=12):
        """计算PSY心理线指标（金融标准）"""
        # 计算上涨天数
        up_days = (df['收盘'] > df['收盘'].shift(1)).astype(int)
        # 明确指定 min_periods，确保使用完整周期数据
        df['PSY'] = up_days.rolling(window=period, min_periods=period).sum() / period * 100
        return df
    
    @staticmethod
    def calculate_all_indicators(df):
        """计算所有技术指标（增强版）"""
        # 原有指标
        df = TechnicalIndicators.calculate_ma(df)
        df = TechnicalIndicators.calculate_macd(df)
        df = TechnicalIndicators.calculate_kdj(df)
        
        # 使用配置的RSI周期
        for period in config.RSI_PERIODS:
            df = TechnicalIndicators.calculate_rsi(df, period)
        
        df = TechnicalIndicators.calculate_boll(df)
        df = TechnicalIndicators.calculate_volume_ma(df)
        
        # 新增指标
        df = TechnicalIndicators.calculate_wr(df, 14)
        df = TechnicalIndicators.calculate_cci(df, 14)
        df = TechnicalIndicators.calculate_atr(df, 14)
        df = TechnicalIndicators.calculate_obv(df)
        df = TechnicalIndicators.calculate_dmi(df, 14)
        df = TechnicalIndicators.calculate_bias(df, [6, 12, 24])
        df = TechnicalIndicators.calculate_roc(df, 12)
        df = TechnicalIndicators.calculate_psy(df, 12)
        
        return df

